import datetime
import pylab as plt
import math
import os
from params import *

STEP = 1
STEP_NAME = 'days'
STEP = 7
STEP_NAME = 'weeks'
#STEP = 30
#STEP_NAME = 'months'
THRESHOLD = 0
MAX = 360/STEP
LIMIT = math.ceil(MAX)
XLBL = 'Activity span in %s'%STEP_NAME

def get_activity(d1,d2):
    format = "%Y-%m-%d"
    d1 = datetime.date(*map(int,d1.split('-')))
    d2 = datetime.date(*map(int,d2.split('-')))
    delta = (d2-d1).days
    delta /= STEP
    return delta
    
def get_age(d1,d2):
    format = "%Y-%m-%d"
    d1 = datetime.date(*map(int,d1.split('-')))
    d2 = datetime.date(2010,8,18)
    delta = (d2-d1).days
    delta /= STEP
    return delta

graph_file = os.path.join(WORKDIR, 'gowalla', 'stats',
    'gowalla_network_snapshot_2010-08-18.txt')
activity_file = os.path.join(WORKDIR, 'gowalla', 'stats',
    'gowalla_user_activity.txt')

user_degree = {}
for line in open(graph_file):
    if line.startswith('#'):
        continue
    n1,n2 = map(int,line.split()[:2])
    user_degree.setdefault(n1,0)
    user_degree.setdefault(n2,0)
    user_degree[n1] += 1
    user_degree[n2] += 1

user_age = {}
user_activity = {}
user_places = {}
user_checkins = {}
for line in open(activity_file):
    user, d1, d2, num_places,num_checkins = line.strip().split()
    user = int(user)
    user_age[user] = get_age(d1,d2)
    user_activity[user] = get_activity(d1,d2)
    user_places[user] = int(num_places)
    user_checkins[user] = int(num_checkins)



def plot(user_values,ylbl,title,mode="LIN",sqrt=False):
    data = {}
    for user in user_values:
        v  = user_values[user]
        #if user in user_activity:
        if user in user_age:
            a = user_activity[user]
            data.setdefault(a,[]).append(v)

    x,y = [],[]
    yerr = []
    for act in sorted(data):
        if THRESHOLD <= act <= MAX:
            x.append(act)
            l = data[act]
            l = map(math.log,l)
            avg = float(sum(l))/len(l)
            var = float(sum((i-avg)**2 for i in l))/len(l)
            std = var**0.5
            avg = math.exp(avg)
            std = math.exp(std)
            if l:
                y.append(avg)
                yerr.append(std)
            else:
                y.append(0.0)
                yerr.append(0.0)
            print x[-1], y[-1], yerr[-1]

    xfit = list(x)
    yfit = list(y)
    y2fit = list(yerr)

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    fmt = 'c+'
    fmt2 = 'co'
    if mode == 'EXP':
        yfit = map(math.log,yfit)
        y2fit = map(math.log,y2fit)
    if sqrt:
        xfit = map(math.sqrt,xfit)

    p = plt.polyfit(xfit,yfit,1)
    p2 = plt.polyfit(xfit,y2fit,1)
    print p
    f = plt.polyval(p,xfit)
    f2 = plt.polyval(p2,xfit)
    if sqrt:
        xfit = map(lambda i: i*i,xfit)
    if mode == "EXP":
        f = map(math.exp,f)
        f2 = map(math.exp,f2)
        p1 = plt.semilogy(x,y,fmt, mfc='None')
        plt.semilogy(x,f,'k--',linewidth=2)
        p2 = plt.semilogy(x,yerr,fmt2, mfc='None', mec='c')
        #plt.semilogy(x,f2,'k--')
    else:
        plt.plot(x,y,fmt, mfc='None')
        plt.plot(x,f,'k--')
        plt.plot(x,yerr,fmt2, mfc='None', mec='c')
    x1,x2,y1,y2 = plt.axis()
    plt.axis((0,LIMIT,1,y2))
    plt.legend([p1[0],p2[0]],
            ['Average', 'Standard deviation'],
            loc = 'upper left',
            numpoints=1)
    plt.ylabel('%s'%ylbl)
    plt.xlabel(XLBL)
    plt.grid(True)
    plt.savefig('gowalla_activity_%s.pdf'%title)
    plt.close()

mode = 'EXP'
plot(user_places,'Places','places',mode=mode,sqrt=True)
plot(user_checkins,'Check-ins','checkins',mode=mode,sqrt=True)
plot(user_degree,'Friends','friends',mode=mode,sqrt=False)
